
import sys

DBSNP_FILE = "/archive10/bowtie/references/known_snps/snp130-singles.txt"

def readDbsnpSnps(filename=DBSNP_FILE):
    dbsnp = {}
    f = open(filename, "r")
    for line in f:
        cols = line.strip().split("\t")
        if len(cols)!=3:
            print >> sys.stderr, "# Warning: ignoring line", line.strip()
            continue
        key = (cols[0], cols[1], cols[2])
        dbsnp[key] = True
        if 0==len(dbsnp)%100000:
            print >> sys.stderr, "# Read %d lines from %s" % (len(dbsnp), filename)
    f.close()
    return dbsnp

def readCountsByClass(filename):
    cts = {}
    f = open(filename, "r")
    for line in f:
        if line.startswith("#"): continue
        cols = line.strip().split("\t")
        key = (cols[0], cols[1])
        cts[key] = "\t".join(cols[2:])
        if 0==len(cts)%100000:
            print >> sys.stderr, "# Read %d lines from %s" % (len(cts), filename)
    f.close()
    return cts

if __name__=="__main__":
    cts = readCountsByClass(sys.argv[1])
    datafile = open(sys.argv[2])
    datafile.close()
    dbsnp = readDbsnpSnps()
    for filename in sys.argv[2:]:
        print >> sys.stderr, "# Processing file", filename
        datafile = open(filename)
        outfile = open(filename.replace(".txt", ".with-sample-classes.txt"), "w")
        for line in datafile:
            cols = line.strip().split("\t")
            key = (cols[0], cols[1], cols[2])
            inDbsnp = "dbsnp-no"
            if dbsnp.has_key(key):
                inDbsnp = "dbsnp-yes"
            outfile.write("%s\t%s\t%s\n" % (line.strip(), inDbsnp, cts[(cols[0], cols[2])]))
        outfile.close()
        datafile.close()
